import torch
from torch import nn
from pfnn_param import PfnnParameter

class PFNN(nn.Module):
    def __init__(self,
                 Xdim,
                 Ydim,
                 num_experts = 4,
                 dropout = .3):
        super().__init__()
        self.num_experts = num_experts
        self.layer1 = PfnnParameter((num_experts, 512, Xdim))
        self.layer2 = PfnnParameter((num_experts, 512, 512))
        self.layer3 = PfnnParameter((num_experts, Ydim, 512))
        self.act = nn.ELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, phase):
        x = self.dropout(x)
        x = self.layer1(x, phase)
        x = self.act(x)
        x = self.dropout(x)
        x = self.layer2(x, phase)
        x = self.act(x)
        x = self.dropout(x)
        x = self.layer3(x, phase)
        x = self.act(x)
        return x